# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import collections
import datetime
import os
import logging
import json

from absl import (flags, app)
import numpy as np

import coin.flow.subscriber as flow_subscriber
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.base.order_gateway import OrderSide, OrderType, Order
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.base.param_util import to_list
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.accounting.pnl_balance_printer.util import reserve_as_proto
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import \
  convert_executor_config_from_dict as convert_pass_config
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.huobi_futures.kr_rest.native_private_client import HuobiFuturesOrderSide
from coin.exchange.okex_v3.kr_rest.native_private_client import OkexFuturesOrderSide
from coin.strategy.tool.fill_log_dumper import FillLogDumper
from coin.strategy.basis_strat.fill_info_writer import FillInfo, FillInfoWriter
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.strategy.accounting.fee import get_fee_rate
from coin.strategy.basis_strat.moving_average import MovingAverage, MovingObject
from coin.support.telemetry.latency_recorder import LatencyRecorder, LifeOfSignal
from coin.strategy.basis_strat.base_strategy import BaseStrategy

import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.fastfeature.sim.lm_sim_executor as lmsimexec
import coin.strategy.mm.fastfeature.sim.lm_strategy_sim as lmssim
import coin.strategy.executor.passive.util as putil

FLAGS = flags.FLAGS


class DummyPosition(object):
  def __init__(self):
    self.long = self.short = 0


def _gen_arb_products(common, products0, products1, ProductType0, ProductType1, start_time):
  if 'exch0_quote_currency' in common and 'exch1_quote_currency' in common:
    exch0_quote_currency = common['exch0_quote_currency']
    exch1_quote_currency = common['exch1_quote_currency']
  else:
    raise ValueError(common)

  products_list0 = []
  products_list1 = []
  for p0, p1 in zip(sorted(list(products0.keys())), sorted(list(products1.keys()))):
    product0 = ProductType0.FromStr(p0, start_time)
    product1 = ProductType1.FromStr(p1, start_time)
    products_list0.append(product0)
    products_list1.append(product1)
  return products_list0, products_list1


import coin.strategy.mm.fastfeature.feed_manager as fmgr
import coin.strategy.mm.fastfeature.fast_feature_feed_spec as fspec
import coin.strategy.mm.fastfeature.feature_feed_converter as ffcv
import coin.strategy.mm.fastfeature.linear_model_factory as lmf
import coin.strategy.mm.fastfeature.util as ffutil
import pyfastfeature.feature.simulator as pysim
from coin.support.latency.latency_model import LatencyModel


class BasisSmmSubStrategy(object):
  ExecutorType = lmsimexec.LmSimMMExecutor
  ExecutorConfigConverter = putil.convert_executor_config_from_dict
  PostOnly = False

  def __init__(self,
               parent,
               config,
               product0,
               product1,
               latency_ms,
               delay_accept_until_book,
               start_time):
    self._product0 = product0
    self._product1 = product1
    assert self._product0.base.currency == self._product1.base.currency
    self._ticker0 = product0.full_symbol
    self._ticker1 = product1.full_symbol
    self._tick_info1 = (get_holder_from_product(product1).product_info.gen_tick(),
                        get_holder_from_product(product1).product_info.gen_tick_qty())
    if hasattr(self._tick_info1[0], 'tick_size'):
      ticksize = self._tick_info1[0].tick_size
      self._next_price_func = lambda p: p + ticksize - (p % ticksize)
      self._prev_price_func = lambda p: p - ((p % ticksize) or ticksize)
    else:
      self._next_price_func = self._tick_info1[0].get_next_price
      self._prev_price_func = self._tick_info1[0].get_prev_price

    self._raw_config1 = None
    self._config1 = None
    self._config1_proto = None
    self.og1 = None
    self._pass_executor1 = None
    self._pass_buy_price = np.nan
    self._pass_sell_price = np.nan

    self._parent = parent
    self.logger = parent.logger
    self.feed = parent._last_feed

    # Params
    self._enter_position_threshold = FLAGS.enter_position_threshold
    self._window_period = np.nan
    self._trading_threshold = FLAGS.trading_threshold_bp * 1e-4
    self._max_std_ratio = FLAGS.max_std_ratio
    self._min_std_ratio = FLAGS.min_std_ratio
    self._price_pull = FLAGS.price_pull_edge * 1e-4
    self._price_push = FLAGS.price_push_edge * 1e-4
    self._max_lean_size = FLAGS.max_lean_bp * 1e-4
    self._pos_lean_mul = FLAGS.pos_lean_multiplier
    self._full_pos_lean_mul = FLAGS.full_pos_lean_multiplier
    self._pos_ratio_func = None
    self._thrift_api = FLAGS.thrift_api
    self._thrift_ignore_levels = FLAGS.thrift_ignore_levels

    self._ma = None
    self._mean = None
    self._std = None
    self._constrained_std = None
    self._book_ts = [0, 0]
    self._book_invalid_period = FLAGS.book_invalid_sec * 10**9

    self._pos1 = 0
    self._pos1_ts = 0
    self._og_pos1 = 0
    self._og_pos1_ts = 0
    self._update_pos_with_balance = FLAGS.update_pos_with_balance
    self._update_position_duration = FLAGS.update_pos_with_balance_sec * 10**9
    self._update_working_qty_duration = FLAGS.update_working_qty_sec * 10**9
    self._equal_theshold = None

    self._lot_size = None
    self._max_pos = None
    self._min_pos = None
    self._agg_taker_fee = None
    self._expected_working_buy = 0
    self._expected_working_sell = 0
    self._last_submit_fill_ts = 0
    # Handle agg working order update delay.
    self._agg_working_order_fills = dict()

    # Decide if book has changed.
    self._book_change_levels = min(FLAGS.book_change_levels, FLAGS.book_levels)
    self._last_book0 = None
    self._last_book1 = None

    self._get_pos_func = None
    self._is_buy_order_func = None
    self._dummy_position = DummyPosition()
    self._raw_run_size = 1.0
    self._relative_run_size = None

    self._force_fire_threshold = FLAGS.force_fire_threshold_bp * 1e-4
    self._force_fire_cooldown = FLAGS.force_fire_cooldown_sec * 1e9
    self._force_fire_cooldown_window = FLAGS.force_fire_cooldown_window_sec * 1e9
    self._force_fire_max_count = FLAGS.force_fire_max_count
    self._force_fire_qty_mul = FLAGS.force_fire_qty_multiplier
    self._force_fire_cap_qty = self._lot_size
    self._agg_buys = MovingObject(self._force_fire_cooldown_window)
    self._agg_sells = MovingObject(self._force_fire_cooldown_window)

    psub_map = {}
    for feed_product in [product0, product1]:
      sub_key = sub.get_subscription_key(feed_product,
                                         api_version_override_map={('Futures', 'Okex'): 'v3'})
      psub_map[feed_product] = sub_key.to_str()

    self._product_map = ffutil.product_map_from_products(list(psub_map.keys()))
    self._feed_products = list(self._product_map.keys())

    self._feedspec = fspec.FastFeatureFeedInfo(
        psub_map=psub_map,
        level_override_map=collections.defaultdict(lambda: 30),
        sub_key_not_none=False,
        bitmex_bbo_compress=False)

    ffcv_converter = ffcv.FeatureFeedCacheConverter if FLAGS.use_feed_cache else ffcv.FeatureArchiveFeedConverter
    self._feed_converter = lmf.create_feed_converter(ffcv_converter,
                                                     self._product_map,
                                                     self._feedspec,
                                                     start_time)

    self._simulator = pysim.Simulator(list(self._product_map.values()),
                                      latency_ns=int(latency_ms * 1e6),
                                      delay_accept_until_book=delay_accept_until_book)

    trade_product = product1

    self._feed_converter.register_symbols(self._simulator)
    self._trade_pindex = self._feed_converter.get_index(trade_product)
    self._logger = parent.logger

    order_gateway = lmsimexec.LmSimOrderGateway(self._logger,
                                                self._simulator,
                                                trade_product,
                                                self._trade_pindex,
                                                post_only=self.PostOnly)
    order_gateway.register_fill_callback(self.on_fill)

    self._execution_params = type(self).ExecutorConfigConverter(trade_product.symbol, config)

    self.init_executors(parent._exch1_subreq, config, order_gateway)

  def _set_pos_lean_ratio_func(self):
    if self._max_lean_size == 0:
      self._pos_ratio_func = lambda _: 0.0
    else:
      pos_range = self._max_pos - self._min_pos
      self._pos_ratio_func = lambda pos: \
          (1 - 2.0 * (pos - self._min_pos) / pos_range)

  def get_states(self):
    return {
        "mean": self._mean,
        "std": self._std,
        "current_mean": self._ma.get_latest_data() if self._ma else np.nan,
        "dq_size": self._ma.size() if self._ma else 0
    }

  def _get_deque_file_name(self):
    return '%s_%s_%s' % \
        (self._product0.exchange, self._product1.exchange, self._ticker0)

  def _init_get_pos_func(self, og):
    product_type = self._product1.product_type
    exchange = self._product1.exchange
    if product_type == 'Futures' or product_type == 'CurrencyPair':
      self._get_pos_func = lambda: og.get_position(self._product1)
      self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    assert self._get_pos_func is not None, '%s %s' % (product_type, exchange)

  def _update_config_from_config_file(self, config):
    config_file_name = FLAGS.config_file_name
    if not config_file_name.endswith('json'):
      config_file_name = '%s.json' % config_file_name
    config_file_path = os.path.join(FLAGS.config_file_dir, config_file_name)
    if not os.path.exists(config_file_path):
      return config
    with open(config_file_path, 'r') as fi:
      data = json.load(fi)
      ticker = self._product1.base.symbol
      if 'params' in data and ticker in data['params']:
        for key, val in data['params'][ticker].items():
          config[key] = float(val)
    self.logger.info('Updated config %s: %s' % (self._ticker1, config))
    return config

  def init_executors(self, req1, config1, og1):
    self._relative_run_size = 1.0
    self.og1 = og1
    if FLAGS.read_config_from_file:
      config1 = self._update_config_from_config_file(config1)
    self._raw_config1 = config1.copy()
    self._config1 = config1.copy()
    self._config1['force_fire'] = True

    if FLAGS.run_size != 1.0:
      assert 0 < FLAGS.run_size < 1, FLAGS.run_size
      self._raw_run_size = FLAGS.run_size
      max_pos = self._config1['max_pos']
      min_pos = self._config1['min_pos']
      neutral_pos = (max_pos + min_pos) / 2.0
      max_exposure = (max_pos - min_pos) / 2.0 * FLAGS.run_size
      self._config1['max_pos'] = neutral_pos + max_exposure
      self._config1['min_pos'] = neutral_pos - max_exposure
      self._config1['lot_size'] = self._config1['lot_size'] * FLAGS.run_size
      self.logger.info('Trimmed size %f %f' % (self._config1['min_pos'], self._config1['max_pos']))

    self._config1_proto = convert_pass_config(self._product1.symbol, self._config1)
    self._pass_executor1 = self.ExecutorType(self._product1,
                                             self._config1_proto,
                                             self.og1,
                                             self.logger)

    self._init_get_pos_func(og1)
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._force_fire_cap_qty = self._lot_size * self._force_fire_qty_mul
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']
    self._max_std_ratio = self._config1.get('max_std_ratio', self._max_std_ratio)
    self._min_std_ratio = self._config1.get('min_std_ratio', self._min_std_ratio)
    self._window_period = self._config1.get('window_period_sec', FLAGS.window_period_sec) * 1e9
    self._pos_lean_mul = self._config1.get('pos_lean_multiplier', self._pos_lean_mul)
    self._full_pos_lean_mul = self._config1.get('full_pos_lean_multiplier', self._full_pos_lean_mul)
    if 'force_fire_threshold_bp' in self._config1:
      self._force_fire_threshold = \
          self._config1['force_fire_threshold_bp'] * 1e-4
    if 'price_pull_edge' in self._config1:
      self._price_pull = self._config1['price_pull_edge'] * 1e-4
    if 'price_push_edge' in self._config1:
      self._price_push = self._config1['price_push_edge'] * 1e-4
    assert 0 < self._min_std_ratio < 1e-2, self._min_std_ratio
    assert 0 < self._max_std_ratio < 1e-2, self._max_std_ratio
    assert self._window_period > 1e9, self._window_period
    assert self._pos_lean_mul >= 0, self._pos_lean_mul
    assert self._full_pos_lean_mul >= 0, self._full_pos_lean_mul
    assert -0.1 < self._force_fire_threshold < 0.1, self._force_fire_threshold
    self._ma = MovingAverage(self._window_period)
    self._enter_position_threshold = self._config1.get('enter_position_threshold',
                                                       self._enter_position_threshold)

    _, self._agg_taker_fee = get_fee_rate(self.og1.market_type, self.og1.exchange, self._product1)
    if hasattr(self._product1, 'expiry') and self._product1.expiry:
      self.logger.info('Futures product, set fees to 0, %s' % self._product1.expiry)
      self._agg_taker_fee = 0.0
    self._update_pos1(self._get_pos_func())
    self._og_pos1 = self._pos1
    self._og_pos1_ts = self._get_ts()
    assert abs(self._agg_taker_fee) <= 1e-2, self._agg_taker_fee
    self.logger.info('%s init pos1 %f' % (self._ticker1, self._pos1))
    self._set_pos_lean_ratio_func()
    assert self._pos_ratio_func is not None

  def _update_pos1(self, pos):
    self._pos1 = pos
    self._pos1_ts = self._get_ts()

  def _get_ts(self):
    return self._parent._ts

  def _is_book_valid(self):
    return self._get_ts() - self._book_ts[0] < self._book_invalid_period

  def _update_deque(self, implied_rate, mid0, mid1):
    ts = self._get_ts()
    self._ma.update(ts, implied_rate)

  def get_implied_rate(self):
    return self._ma.get_latest_data()

  def _update_statistics(self):
    self._mean, self._std = self._ma.get_statistics()
    self._constrained_std = self._std
    self._constrained_std = min(self._constrained_std, self._max_std_ratio * self._mean)
    self._constrained_std = max(self._constrained_std, self._min_std_ratio * self._mean)

  def ready(self):
    return self._ma.ready()

  def _position_equal(self, p1, p2):
    return abs(p1 - p2) < self._equal_theshold

  def _get_working_qty(self):
    # Use working orders to update |self._expected_working_qty|.
    if self._get_ts() - self._last_submit_fill_ts > self._update_working_qty_duration:
      order_id_set = set()
      working_buy = 0.0
      working_sell = 0.0
      for order in self._pass_executor1._wo.buy_working_orders:
        working_buy += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      for order in self._pass_executor1._wo.sell_working_orders:
        working_sell += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      # Remove order ids that are not in working orders.
      obsolete_ids = self._agg_working_order_fills.keys() - order_id_set
      for order_id in obsolete_ids:
        del self._agg_working_order_fills[order_id]
      if not self._position_equal(self._expected_working_buy, working_buy):
        self.logger.info('working qty expected %f, updated %f' %
                         (self._expected_working_buy, working_buy))
      if not self._position_equal(self._expected_working_sell, working_sell):
        self.logger.info('working qty expected %f, updated %f' %
                         (self._expected_working_sell, working_sell))
      self._expected_working_buy = working_buy
      self._expected_working_sell = working_sell
    return self._expected_working_buy, self._expected_working_sell

  def _is_pos_valid(self):
    ts = self._get_ts()
    return ts - max(self._pos1_ts or 0, self._og_pos1_ts) > self._update_position_duration

  def _get_balance(self):
    og_pos = self._get_pos_func()
    if self._update_pos_with_balance:
      if not self._position_equal(og_pos, self._og_pos1):
        self._og_pos1 = og_pos
        self._og_pos1_ts = self._get_ts()
      if not self._position_equal(self._pos1, self._og_pos1) \
          and self._is_pos_valid():
        self.logger.info('Update pos1 with og_pos1, %f %f' % (self._pos1, self._og_pos1))
        self._pos1 = self._og_pos1
    else:
      if self._position_equal(self._pos1, og_pos):
        self._pos1 = og_pos
    return self._pos1

  def _get_agg_working_qty(self, mo):
    qty = 0.0
    for ts, obj in mo.iteritems():
      qty += obj[3]
    return qty

  def _get_force_buy_qty(self, raw_price, agg_price, asks):
    self._agg_buys.update(self._get_ts())
    size = self._agg_buys.size()
    qty = 0
    working_qty = 0
    avail_qty = 0
    if size >= self._force_fire_max_count:
      return 0
    elif size > 0:
      last_ts = self._agg_buys.get_latest_ts()
      if self._get_ts() - last_ts < self._force_fire_cooldown:
        return 0
      last_raw_price, last_agg_price, last_ba, _ = \
          self._agg_buys.get_latest_data()
      # Order book not updated and agg price not improved.
      if np.isclose(last_ba[0], asks[0][0]) \
          and np.isclose(last_ba[1], asks[0][1]) \
          and agg_price <= last_agg_price + 1e-6:
        return 0
      # Agg raw price moving lower, do not fire.
      if raw_price < last_raw_price:
        return 0
      working_qty = self._get_agg_working_qty(self._agg_buys)
    for p, q in asks:
      if p > raw_price:
        break
      avail_qty += q
    vacancy_qty = self._max_pos - self._get_balance() - working_qty
    qty = min(vacancy_qty, self._force_fire_cap_qty, avail_qty)
    return qty

  def _get_force_sell_qty(self, raw_price, agg_price, bids):
    self._agg_sells.update(self._get_ts())
    size = self._agg_sells.size()
    qty = 0
    working_qty = 0
    avail_qty = 0
    if size >= self._force_fire_max_count:
      return 0
    elif size > 0:
      last_ts = self._agg_sells.get_latest_ts()
      if self._get_ts() - last_ts < self._force_fire_cooldown:
        return 0
      last_raw_price, last_agg_price, last_bb, _ = \
          self._agg_sells.get_latest_data()
      # Order book not updated and agg price not improved.
      if np.isclose(last_bb[0], bids[0][0]) \
          and np.isclose(last_bb[1], bids[0][1]) \
          and agg_price >= last_agg_price - 1e-6:
        return 0
      # Agg raw price moving higher, do not fire.
      if raw_price > last_raw_price:
        return 0
      working_qty = self._get_agg_working_qty(self._agg_sells)
    for b, q in bids:
      if b < raw_price:
        break
      avail_qty += q
    vacancy_qty = self._get_balance() - self._min_pos - working_qty
    qty = min(vacancy_qty, self._force_fire_cap_qty, avail_qty)
    return qty

  def _process(self, implied_rate, bids0, asks0, bids1, asks1, mid0, mid1):
    self._update_statistics()
    shift_std = self._constrained_std * self._enter_position_threshold
    shift_edge = shift_std / self._mean
    ref_price = self._parent.get_mid(bids0, asks0)
    pos1 = self._get_balance()
    pos_lean_ratio = self._pos_ratio_func(pos1)
    lean_size = self._pos_lean_mul * shift_edge * pos_lean_ratio
    capped_lean_size = max(min(lean_size, self._max_lean_size), -self._max_lean_size)
    if lean_size > 0:
      # Negative position
      raw_pass_buy_price = ref_price * \
          (1 + capped_lean_size) * (self._mean - shift_std)
      raw_pass_sell_price = ref_price * \
          (1 + lean_size * self._full_pos_lean_mul) * (self._mean + shift_std)
    else:
      # Positive position
      raw_pass_buy_price = ref_price * \
          (1 + lean_size * self._full_pos_lean_mul) * (self._mean - shift_std)
      raw_pass_sell_price = ref_price * \
          (1 + capped_lean_size) * (self._mean + shift_std)

    if raw_pass_buy_price * (1 - self._force_fire_threshold) > asks1[0][0]:
      agg_buy_price = float(self._prev_price_func(raw_pass_buy_price))
      agg_buy_qty = self._get_force_buy_qty(raw_pass_buy_price, agg_buy_price, asks1)
      if not FLAGS.dry_run and agg_buy_qty > 0:
        self._pass_executor1.force_fire_buy_order(agg_buy_price, agg_buy_qty)
        self._agg_buys.update(self._get_ts(),
                              (raw_pass_buy_price, agg_buy_price, asks1[0], agg_buy_qty))
        # self._expected_working_buy += agg_buy_qty
        # self.logger.info('Force buy %s %f @ %f' % (self._ticker1, agg_buy_price, agg_buy_qty))
    if raw_pass_sell_price * (1 + self._force_fire_threshold) < bids1[0][0]:
      agg_sell_price = float(self._next_price_func(raw_pass_sell_price))
      agg_sell_qty = self._get_force_sell_qty(raw_pass_sell_price, agg_sell_price, bids1)
      if not FLAGS.dry_run and agg_sell_qty > 0:
        self._pass_executor1.force_fire_sell_order(agg_sell_price, agg_sell_qty)
        self._agg_sells.update(self._get_ts(),
                               (raw_pass_sell_price, agg_sell_price, bids1[0], agg_sell_qty))
        # self._expected_working_sell += agg_sell_qty
        # self.logger.info('Force sell %s %f @ %f' % (self._ticker1, agg_sell_price, agg_sell_qty))

    self._pass_buy_price = self._prev_price_func(min(raw_pass_buy_price, asks1[0][0] - 1e-6))
    self._pass_sell_price = self._next_price_func(max(raw_pass_sell_price, bids1[0][0] + 1e-6))
    buy_price_pull = self._pass_buy_price * (1 + self._price_pull)
    sell_price_pull = self._pass_sell_price * (1 - self._price_pull)
    buy_price_push = self._pass_buy_price * (1 - self._price_push)
    sell_price_push = self._pass_sell_price * (1 + self._price_push)
    if self._thrift_api:
      if self._pass_buy_price < bids1[self._thrift_ignore_levels - 1][0]:
        self._pass_buy_price = buy_price_push = None
      if self._pass_sell_price > asks1[self._thrift_ignore_levels - 1][0]:
        self._pass_sell_price = sell_price_push = None
    if not FLAGS.dry_run:
      self._pass_executor1.set_timestamp(self._get_ts())
      # self._pass_executor1._update_state(self._get_ts())
      self._pass_executor1.manage_mm_orders(sell_price=self._pass_sell_price,
                                            buy_price=self._pass_buy_price,
                                            sell_price_pull=sell_price_pull,
                                            buy_price_pull=buy_price_pull,
                                            sell_price_push=sell_price_push,
                                            buy_price_push=buy_price_push)

  def on_book_record_only(self, exchange_id, book):
    self._book_ts[exchange_id] = book.timestamp
    if not self._is_book_valid():
      # self.logger.warning('%s book is not updated %s %s.' % (self._ticker0, self._book_ts[0], self._book_ts[1]))
      return
    bids0, asks0 = self.feed[self._product0]
    bids1, asks1 = self.feed[self._product1]
    book_changed = self._book_changed(bids0, asks0, self._last_book0) \
        if exchange_id == 0 else self._book_changed(bids1, asks1, self._last_book1)
    if book_changed:
      mid0 = self._parent.get_mid(bids0, asks0)
      mid1 = self._parent.get_mid(bids1, asks1)
      implied_rate = mid1 / mid0
      self._update_deque(implied_rate, mid0, mid1)
    if exchange_id == 0:
      self._last_book0 = (bids0, asks0)
    else:
      self._last_book1 = (bids1, asks1)

  def output_deque(self):
    file_name = '%s_%s_%s' % \
        (self._product0.exchange, self._product1.exchange, self._ticker0)
    self.logger.info('Start writing deque %s' % file_name)
    with open(file_name, 'w') as fo:
      content = '\n'.join(['%d,%f' % (e[0], e[1]) for e in self._ma._dq])
      fo.write(content)
    self.logger.info('Finish writing deque')

  def _book_changed(self, bids, asks, last_book):
    if last_book:
      last_bids, last_asks = last_book
      for i in range(min(self._book_change_levels, len(bids))):
        try:
          if tuple(bids[i]) != tuple(last_bids[i]) or tuple(asks[i]) != tuple(last_asks[i]):
            return True
        except Exception:
          import pdb
          pdb.set_trace()
    return False

  def on_book(self, exchange_id, product, book):
    if not self._feed_converter.update_books(product, book, [self._simulator]):
      return
    self._book_ts[exchange_id] = book.timestamp
    if self._pass_executor1:
      if not self._is_book_valid():
        # self.logger.warning('%s book is not updated %s %s.' % (self._ticker0, self._book_ts[0], self._book_ts[1]))
        return

      if self._product0 not in self.feed or self._product1 not in self.feed:
        return

      bids0, asks0 = self.feed[self._product0]
      bids1, asks1 = self.feed[self._product1]
      book_changed = self._book_changed(bids0, asks0, self._last_book0) \
          if exchange_id == 0 else self._book_changed(bids1, asks1, self._last_book1)
      if book_changed:
        mid0 = self._parent.get_mid(bids0, asks0)
        mid1 = self._parent.get_mid(bids1, asks1)
        implied_rate = mid1 / mid0
        self._update_deque(implied_rate, mid0, mid1)

        if self.ready():
          self._process(implied_rate, bids0, asks0, bids1, asks1, mid0, mid1)
      if exchange_id == 0:
        self._last_book0 = (bids0, asks0)
      else:
        self._last_book1 = (bids1, asks1)

  def on_fill(self, fill_price, fill_qty, order):
    self._agg_working_order_fills[order.order_id] = \
        self._agg_working_order_fills.get(order.order_id, 0.0) + fill_qty
    if self._is_buy_order_func(order.order_side):
      self._update_pos1(self._pos1 + fill_qty * (1 - self._agg_taker_fee))
      # self._expected_working_buy -= fill_qty
    else:
      self._update_pos1(self._pos1 - fill_qty)
      # self._expected_working_sell -= fill_qty
    self._last_submit_fill_ts = self._get_ts()
    self.logger.info('expected %s pos, %f' % (self._ticker1, self._pos1))

  def on_trade(self, exchange_id, product, trade):
    if not self._feed_converter.update_trade(product, trade, self._simulator):
      return


class BasisSmmSimStrategy(lmssim.LmSimStrategy):
  def __init__(self,
               strategy_params,
               trading_tickers,
               record_implied_rate_only=False,
               latency_ms=None,
               delay_accept_until_book=None,
               use_predicted_latency=None,
               start_time=None):
    self.logger = logging.getLogger(__name__)
    self._config = strategy_params['config']

    if use_predicted_latency:
      model = LatencyModel()
      path = os.path.join(
          os.path.dirname(os.path.realpath(__file__)),
          '../../../../support/latency/latency.exchange+market_type.20190807-20190815.model')
      model.load(path)
      trade_subreq = self._config['common']['exch1_subreq']
      try:
        latency_ms = int(
            model.predict({
                'exchange': trade_subreq['exchange'], 'market_type': trade_subreq['market_type']
            }) / 1e6)
        print('predicted latency: ',
              latency_ms,
              'ms for ',
              trade_subreq['exchange'],
              trade_subreq['market_type'])
      except Exception:
        print('exception raised while predicting latency')
        print('for ', trade_subreq['exchange'], trade_subreq['market_type'])
    self._trading_tickers = set(trading_tickers)
    self.logger.info('Traded tickers %s' % self._trading_tickers)
    self._record_implied_rate_only = record_implied_rate_only
    self.logger.info('Record implied rate only %s' % self._record_implied_rate_only)

    self._exch0_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch0_subreq']))
    self._exch1_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch1_subreq']))

    self._exchange0 = self._exch0_subreq.exchange
    self._exchange1 = self._exch1_subreq.exchange

    self._book_builder_clses = (get_flow_book_builder_cls(self._exch0_subreq),
                                get_flow_book_builder_cls(self._exch1_subreq))

    self._live_env_clses1 = get_live_env_cls(self._exch1_subreq)
    self._live_env1 = None
    self.og1 = None

    self._products0, self._products1 = _gen_arb_products(self._config['common'],
                                                         self._config['products0'],
                                                         self._config['products1'],
                                                         self._book_builder_clses[0].ProductType,
                                                         self._book_builder_clses[1].ProductType,
                                                         start_time)
    self._is_future = self._products1[0].product_type == 'Futures'

    exch0_exchange_feed_checker = self._config['common'].get('exch0_exchange_feed_checker', False)
    exch1_exchange_feed_checker = self._config['common'].get('exch1_exchange_feed_checker', False)
    feed_checker_products = []
    if exch0_exchange_feed_checker:
      feed_checker_products += self._products0
    if exch1_exchange_feed_checker:
      feed_checker_products += self._products1
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    self._feed = MultiFeedSubsystem(feed_checker_config=feed_checker_config)

    # Shrink size when rollover is close.
    self._real_time = datetime.datetime.now()
    expiry = self._get_config_key(self._products0[0])
    self._ref_weekly = 'THIS_WEEK' in expiry
    self._ref_quarter = 'QUARTER' in expiry
    self.logger.info('%s, ref weekly: %s, ref quarter: %s' %
                     (expiry, self._ref_weekly, self._ref_quarter))

    self._ts = 0
    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._balance_printer = None
    self._fill_dumper = FillLogDumper(self._exch1_subreq)

    self._book_levels = FLAGS.book_levels
    self._last_feed = {}
    self._book_map = BookMap(self._products0 + self._products1, logger=self.logger)

    self._reserve_map = self._get_reserve_map()
    self._strategy_request = convert_strategy_request_from_dict(
        self._config['common']['strategy_request'])
    self._reserve_proto = reserve_as_proto(self._exch1_subreq, self._reserve_map)

    self._substrategies = {}
    self._strats = {}
    for product0, product1 in zip(self._products0, self._products1):
      key = self._get_config_key(product1)
      print(self._config.keys())
      print(self._config['products1'])
      print(key)
      config1 = self._config['products1'][key].copy()
      config1.update(self._config.get('pass_executor_configs', {}))
      self._substrategies[product0] = self._substrategies[product1] = \
          BasisSmmSubStrategy(
              self, config1, product0, product1, latency_ms, delay_accept_until_book, start_time)

      self._strats[product1] = self._substrategies[product0]

    self._feed_manager = fmgr.FeedManager(self.on_book,
                                          self.on_trade, {},
                                          feedspec=self._substrategies[product1]._feedspec)

  def get_result(self):
    ret_dict = {}
    for trade_product, substrat in self._strats.items():
      ret_dict[trade_product.symbol] = (trade_product,
                                        substrat._simulator.GetHistoricalResponse(
                                            substrat._trade_pindex),
                                        None)
    return ret_dict

  @staticmethod
  def get_mid(bids, asks):
    return (bids[0][0] + asks[0][0]) / 2.0

  def _update_prices(self, product, book):
    bids = book.get_bid_array(self._book_levels)
    asks = book.get_ask_array(self._book_levels)
    self._last_feed[product] = (bids, asks)

  def get_implied_rate(self):
    rates = [strat.get_implied_rate() for strat in self._substrategies.values()]
    return 1.0 / np.nanmean(rates)

  def on_book(self, product, book):
    if product in self._products0:
      exchange_id = 0
    elif product in self._products1:
      exchange_id = 1
    self._ts = book.timestamp
    self._update_prices(product, book)
    self._substrategies[product].on_book(exchange_id, product, book)
    self._book_map.set_book(product, book)

  def on_book_record_only(self, product, book):
    if product in self._products0:
      exchange_id = 0
    elif product in self._products1:
      exchange_id = 1
    self._ts = book.timestamp
    self._update_prices(book)
    self._substrategies[product].on_book_record_only(exchange_id, book)

  def on_trade(self, product, trade):
    if product in self._products0:
      exchange_id = 0
    elif product in self._products1:
      exchange_id = 1
    self._substrategies[product].on_trade(exchange_id, product, trade)

  def on_fill(self, fill_price, fill_qty, order):
    self._substrategies[order.product].on_fill(fill_price, fill_qty, order)

  def on_book_reset(self, book_builder_name, book_builder):
    if self._record_implied_rate_only:
      self.on_book = self.on_book_record_only

    feed_manager = self._feed_manager
    feed_manager.set_book_reset(self._products0, book_builder_name, book_builder)
    feed_manager.set_book_reset(self._products1, book_builder_name, book_builder)

  def _get_reserve_map(self):
    reserves = dict()
    for product in self._products1:
      key = self._get_config_key(product)
      reserves[product.base.currency] = (self._config['products1'][key]['max_pos']
                                         + self._config['products1'][key]['min_pos']) / 2.0
    return reserves

  @staticmethod
  def _get_config_key(product):
    key = product.subscription_symbol
    if product.exchange == 'Huobi' and product.product_type == 'Futures':
      key = key.replace('CURRENT', 'THIS')
      key = key.replace('THIS_QUARTER', 'QUARTER')
    return key


def get_param_set(config_filename):
  config = load_clean_json(config_filename)
  return {
      'config': config,
  }


def get_start_end_time(trading_date, time_range):
  add_hour, hours = time_range.split("-")
  add_hour = float(add_hour)
  hours = float(hours) - add_hour
  start_time = trading_date + datetime.timedelta(hours=add_hour)
  end_time = start_time + datetime.timedelta(hours=hours)
  return start_time, end_time


import coin.proto.archive_flag_pb2 as afpb
import coin.strategy.mm.subscription as sub


def update_flags(**sink_kwargs):
  allflags = set(flags.FLAGS)
  for k, v in sink_kwargs.items():
    if k in allflags and getattr(flags.FLAGS, k) != v:
      setattr(flags.FLAGS, k, v)


def dump_basis_smm_strategy_sim(*,
                                logger,
                                config_filename,
                                trading_date,
                                time_range,
                                latency_ms,
                                use_predicted_latency,
                                **sink_kwargs):
  update_flags(**sink_kwargs)
  if FLAGS.thrift_api:
    assert FLAGS.thrift_ignore_levels <= FLAGS.book_levels
  strategy_params = get_param_set(config_filename=config_filename)

  trading_tickers = FLAGS.trading_tickers.split(',')
  record_implied_rate_only = FLAGS.record_implied_rate_only

  start_time, end_time = get_start_end_time(trading_date, time_range)

  dumpstrategy = BasisSmmSimStrategy(strategy_params,
                                     trading_tickers,
                                     record_implied_rate_only=record_implied_rate_only,
                                     latency_ms=latency_ms,
                                     use_predicted_latency=use_predicted_latency,
                                     start_time=start_time,
                                     delay_accept_until_book=sink_kwargs['delay_accept_until_book'])

  baseflags = afpb.BaseFlagProto(trading_date=trading_date.strftime("%Y%m%d"),
                                 time_range=time_range)

  fsflags = abase.get_feed_subscription_flags()
  fs_req = sub.FeedSubscriptionRequest()
  fs_req.add_products(dumpstrategy._products0, dumpstrategy._exch0_subreq)
  fs_req.add_products(dumpstrategy._products1, dumpstrategy._exch1_subreq)

  abase.run_from_archive_base_with_feed_sub(baseflags=baseflags,
                                            feed_sub_flags_proto=fsflags,
                                            feed_sub_request=fs_req,
                                            on_book_reset=dumpstrategy.on_book_reset)

  return dumpstrategy.get_result()


def define_basis_smm_flags():
  flags.DEFINE_string('trading_tickers', '', '')

  flags.DEFINE_float('book_invalid_sec', 3.5, '')

  flags.DEFINE_boolean('update_pos_with_balance', False, '')

  flags.DEFINE_float('update_pos_with_balance_sec', 4, '')

  flags.DEFINE_float('update_working_qty_sec', 3, '')

  flags.DEFINE_integer('window_period_sec', 600, '')

  flags.DEFINE_float('trading_threshold_bp', 5, '')

  flags.DEFINE_float('enter_position_threshold', 3.0, '')

  flags.DEFINE_boolean('record_implied_rate_only', False, '')

  flags.DEFINE_integer('book_levels', 5, '')

  flags.DEFINE_integer('book_change_levels', 3, '')

  flags.DEFINE_float('max_std_ratio', 10e-4, '')

  flags.DEFINE_float('min_std_ratio', 5e-4, '')

  flags.DEFINE_string('strategy_name', '', '')

  flags.DEFINE_float('price_pull_edge', 2, '')

  flags.DEFINE_float('price_push_edge', 4, '')

  flags.DEFINE_float('max_lean_bp', 10, '')

  flags.DEFINE_float('pos_lean_multiplier', 0.6, '')

  flags.DEFINE_float('full_pos_lean_multiplier', 1.0, '')

  flags.DEFINE_float('force_fire_threshold_bp', 1, '')

  flags.DEFINE_float('force_fire_cooldown_window_sec', 5, '')

  flags.DEFINE_float('force_fire_cooldown_sec', 1, '')

  flags.DEFINE_integer('force_fire_max_count', 3, '')

  flags.DEFINE_float('force_fire_qty_multiplier', 1.2, '')

  flags.DEFINE_boolean('thrift_api', False, '')

  flags.DEFINE_integer('thrift_ignore_levels', 3, '')

  flags.DEFINE_float('run_size', 1.0, '')

  flags.DEFINE_boolean('read_config_from_file', False, '')

  flags.DEFINE_boolean('dry_run', False, '')

  flags.DEFINE_string('config_file_dir', 'coin_deploy/basis_strat/profile', '')

  flags.DEFINE_string('config_file_name', '', '')

  flags.DEFINE_bool(
      'enable_telemetry',
      False,
      'enable telemetry',
  )

  flags.DEFINE_boolean('subscribe_trade', False, '')

  flags.DEFINE_string('maintenance_file', 'coin_deploy/basis_strat/schedule/maintenance.json', '')
